# coding: utf-8
import pickle
import random
import threading
from random import shuffle

import numpy as np
import matplotlib.pyplot as plt

from TSP import TSP

np.set_printoptions(suppress=True, linewidth=999, threshold=np.inf)
plt.rc('font', family='Times New Roman', weight='normal', size=12)


def copy(target):
    return pickle.loads(pickle.dumps(target))


class Encoding:
    def __init__(self, data, cost_func):
        self.cities = None
        self.cost = 1e99
        # init params
        self.cities = copy(data.cities)
        self.depot = data.depot
        self.cities.remove(self.depot)
        shuffle(self.cities)
        self.cities.insert(0, self.depot)
        self.cost_func = cost_func

    def __repr__(self):
        return f'cities: {self.cities}, total: {self.cost}'

    def calc_cost(self):
        self.cost = self.cost_func(route=self.cities + [self.depot])
        return self


class PGA:
    def __init__(self, data):
        self.data = data
        self.cost = data.cost
        self.NP = 100
        self.Population = [Encoding(self.data, self.cost).calc_cost() for _ in range(self.NP)]
        self.gBest = copy(min(self.Population, key=lambda x: x.cost))
        self.evo_list = []
        
    def optimize(self, iterations=1000):
        for it in range(iterations):
            Offspring = []
            for n in range(0, self.NP, 10):
                best = min(self.Population[n: n+10], key=lambda x: x.cost)
                temporary_population = [best for _ in range(10)]
                i, j, p = sorted(random.sample(range(len(best.cities)), k=3))
                i, j, p = (i, j, p) if random.getrandbits(1) else (j, p, i)
                for idx in range(10):
                    if idx == 0:
                        pass
                    elif idx == 1:
                        temporary_population[idx].cities = flip_insert(temporary_population[idx].cities, i, j, p)
                    elif idx == 2:
                        temporary_population[idx].cities = swap_insert(temporary_population[idx].cities, i, j, p)
                    elif idx == 3:
                        temporary_population[idx].cities = left_slide_insert(temporary_population[idx].cities, i, j, p)
                    elif idx == 4:
                        temporary_population[idx].cities = right_slide_insert(temporary_population[idx].cities, i, j, p)
                    elif idx == 5:
                        temporary_population[idx].cities = flip_insert(temporary_population[idx].cities, i, j, p)
                        temporary_population[idx].cities = swap_insert(temporary_population[idx].cities, i, j, p)
                    elif idx == 6:
                        temporary_population[idx].cities = flip_insert(temporary_population[idx].cities, i, j, p)
                        temporary_population[idx].cities = left_slide_insert(temporary_population[idx].cities, i, j, p)
                    elif idx == 7:
                        temporary_population[idx].cities = flip_insert(temporary_population[idx].cities, i, j, p)
                        temporary_population[idx].cities = right_slide_insert(temporary_population[idx].cities, i, j, p)
                    elif idx == 8:
                        temporary_population[idx].cities = swap_insert(temporary_population[idx].cities, i, j, p)
                        temporary_population[idx].cities = left_slide_insert(temporary_population[idx].cities, i, j, p)
                    elif idx == 9:
                        temporary_population[idx].cities = swap_insert(temporary_population[idx].cities, i, j, p)
                        temporary_population[idx].cities = right_slide_insert(temporary_population[idx].cities, i, j, p)
                    temporary_population[idx] = temporary_population[idx].calc_cost()
                Offspring.extend(temporary_population)
            self.Population = copy(Offspring)
            # update gBest
            if self.gBest.cost > min(Offspring, key=lambda x: x.cost).cost:
                self.gBest = copy(min(Offspring, key=lambda x: x.cost))
            self.evo_list.append(self.gBest.cost)
            
            # visualization
            plt.clf()
            plt.subplot(1, 2, 1)
            self.data.draw(plt, routes=[self.gBest.cities])
            plt.title(self.gBest.cost)
            plt.subplot(1, 2, 2)
            plt.plot(range(len(self.evo_list)), self.evo_list, label=f'cost: {self.evo_list[-1]:.2f}')
            plt.title(f'optimal cost')
            plt.pause(0.001)
        plt.show()
                        
        
def flip_insert(route, i, j, p):
    flip = list(reversed(route[i:j]))
    route = route[:i] + route[j:]
    if j < p:
        k = p - j + i
        route = route[:k] + flip + route[k:]
    else:
        route = route[:p] + flip + route[p:]
    return route


def swap_insert(route, i, j, p):
    swap = route[i:j]
    swap[0], swap[-1] = swap[-1], swap[0]
    route = route[:i] + route[j:]
    if j < p:
        k = p - j + i
        route = route[:k] + swap + route[k:]
    else:
        route = route[:p] + swap + route[p:]
    return route


def left_slide_insert(route, i, j, p):
    segment = route[i:j]
    segment = segment[1:] + [segment[0]]
    route = route[:i] + route[j:]
    if j < p:
        k = p - j + i
        route = route[:k] + segment + route[k:]
    else:
        route = route[:p] + segment + route[p:]
    return route
  
    
def right_slide_insert(route, i, j, p):
    segment = route[i:j]
    segment = [segment[-1]] + segment[:-1]
    route = route[:i] + route[j:]
    if j < p:
        k = p - j + i
        route = route[:k] + segment + route[k:]
    else:
        route = route[:p] + segment + route[p:]
    return route
                
    
plt.figure(figsize=(10, 4))
tsp = TSP(path=f'../../data/TSPLIB/burma14.tsp.txt')
optimizer = PGA(data=tsp)
optimizer.optimize()
